import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Circle

class PathPlotter:
    def __init__(self, grid_map, resolution=1.0):
        """
        Initialize path plotter
        :param grid_map: GridMap object
        :param resolution: grid resolution in meters
        """
        self.grid_map = grid_map
        self.resolution = resolution
        self.fig, self.ax = plt.subplots(figsize=(10, 10))
        self.ax.set_aspect('equal')
        self.ax.set_title('A* Path Planning Simulation')
        self.ax.set_xlabel('X (m)')
        self.ax.set_ylabel('Y (m)')

        # Initialize plot with grid map
        self._plot_grid_map()

    def _plot_grid_map(self):
        """Plot the grid map with obstacles"""
        grid = self.grid_map.get_grid()
        height, width = grid.shape

        # Create a color map where obstacles are black and free space is white
        cmap = plt.cm.gray
        cmap.set_bad(color='black')
        grid_plot = np.ma.masked_where(grid == 1, grid)

        # Plot the grid
        self.ax.imshow(grid_plot, cmap=cmap, extent=[0, width*self.resolution, 0, height*self.resolution],
                      origin='lower')

    def plot_path(self, path, color='r', linewidth=2, label='Path'):
        """
        Plot a path on the grid
        :param path: list of [x, y] or [x, y, theta] points in world coordinates
        :param color: path color
        :param linewidth: path line width
        :param label: path label for legend
        """
        if not path:
            print("No path to plot")
            return

        # Extract x and y coordinates
        x = [p[0] for p in path]
        y = [p[1] for p in path]

        # Plot path
        self.ax.plot(x, y, color=color, linewidth=linewidth, label=label)

        # Plot start and end points
        self.ax.plot(x[0], y[0], 'go', markersize=10, label='Start')
        self.ax.plot(x[-1], y[-1], 'bo', markersize=10, label='Goal')

        # Add direction arrows for Hybrid A* path
        if len(path[0]) > 2:
            self._plot_direction_arrows(path)

    def _plot_direction_arrows(self, path, arrow_interval=5, arrow_length=1.0):
        """Plot direction arrows along the path"""
        for i in range(0, len(path), arrow_interval):
            x, y, theta = path[i][:3]
            dx = arrow_length * np.cos(theta)
            dy = arrow_length * np.sin(theta)
            self.ax.arrow(x, y, dx, dy, head_width=0.5, head_length=0.5, fc='g', ec='g')

    def plot_vehicle(self, x, y, theta, color='blue', size=2.5):
        """
        Plot a simple vehicle model
        :param x: x position
        :param y: y position
        :param theta: orientation angle (rad)
        :param color: vehicle color
        :param size: vehicle size scaling factor
        """
        # Vehicle dimensions
        length = 4.0 * size
        width = 1.8 * size

        # Vehicle corners
        corners = [
            [length/2, width/2],
            [length/2, -width/2],
            [-length/2, -width/2],
            [-length/2, width/2]
        ]

        # Rotate and translate corners
        rotated_corners = []
        for (cx, cy) in corners:
            rx = cx * np.cos(theta) - cy * np.sin(theta) + x
            ry = cx * np.sin(theta) + cy * np.cos(theta) + y
            rotated_corners.append([rx, ry])

        # Close the polygon
        rotated_corners.append(rotated_corners[0])

        # Plot vehicle
        x_vals = [p[0] for p in rotated_corners]
        y_vals = [p[1] for p in rotated_corners]
        self.ax.plot(x_vals, y_vals, color=color)

        # Plot front direction
        front_x = x + length/2 * np.cos(theta)
        front_y = y + length/2 * np.sin(theta)
        self.ax.plot([x, front_x], [y, front_y], 'r-', linewidth=2)

    def plot_obstacles(self, obstacles, color='k', size=0.5):
        """
        Plot additional obstacles
        :param obstacles: list of (x, y) obstacle positions
        :param color: obstacle color
        :param size: obstacle size
        """
        for (x, y) in obstacles:
            self.ax.add_patch(Circle((x, y), size, color=color))

    def plot_dynamic_obstacles(self, obstacles, color='orange', size=0.7):
        """
        Plot dynamic obstacles
        :param obstacles: list of (x, y) obstacle positions
        :param color: obstacle color
        :param size: obstacle size
        """
        for (x, y) in obstacles:
            self.ax.add_patch(Circle((x, y), size, color=color, alpha=0.7))
            self.ax.text(x, y, 'D', ha='center', va='center', color='white')

    def show(self, block=True):
        """Display the plot"""
        self.ax.legend()
        plt.tight_layout()
        plt.show(block=block)

    def clear(self):
        """Clear all plotted elements except the grid map"""
        self.ax.clear()
        self._plot_grid_map()

    def save_figure(self, filename='path_planning_result.png'):
        """Save the current plot to a file"""
        self.ax.legend()
        plt.tight_layout()
        plt.savefig(filename, dpi=300)
        print(f"Figure saved as {filename}")